#!/bin/bash

# YOLOv8 OpenVINO Go Service Client
# Shell script for interacting with the inference service

set -e

# Configuration
SERVICE_URL="${SERVICE_URL:-http://localhost:8080}"
TIMEOUT=30

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# Functions
print_usage() {
    cat << EOF
Usage: $0 [COMMAND] [OPTIONS]

Commands:
  health              Check service health
  info                Get model information
  infer IMAGE_PATH    Run inference on single image
  batch IMAGE1 IMAGE2 ... Run batch inference on multiple images

Options:
  --url URL           Service URL (default: http://localhost:8080)
  --timeout SECONDS   Request timeout (default: 30)
  --help              Show this help message

Examples:
  $0 health
  $0 --url http://example.com:8080 info
  $0 infer /path/to/image.jpg
  $0 batch /path/to/image1.jpg /path/to/image2.jpg
EOF
}

health_check() {
    echo -e "${YELLOW}Checking service health...${NC}"
    response=$(curl -s -m "$TIMEOUT" "$SERVICE_URL/health")
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}Service is healthy${NC}"
        echo "$response" | jq '.'
        return 0
    else
        echo -e "${RED}Service health check failed${NC}"
        return 1
    fi
}

get_model_info() {
    echo -e "${YELLOW}Getting model information...${NC}"
    response=$(curl -s -m "$TIMEOUT" "$SERVICE_URL/api/v1/model/info")
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}Model information:${NC}"
        echo "$response" | jq '.'
        return 0
    else
        echo -e "${RED}Failed to get model information${NC}"
        return 1
    fi
}

infer_single() {
    local image_path=$1
    
    if [ ! -f "$image_path" ]; then
        echo -e "${RED}Error: Image not found: $image_path${NC}"
        return 1
    fi
    
    echo -e "${YELLOW}Running inference on $image_path...${NC}"
    
    response=$(curl -s -m "$TIMEOUT" -X POST "$SERVICE_URL/api/v1/infer" \
        -H "Content-Type: application/json" \
        -d "{\"image_path\": \"$image_path\"}")
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}Inference completed${NC}"
        echo "$response" | jq '.'
        
        # Extract and display detections
        detections=$(echo "$response" | jq '.detections | length')
        echo -e "\n${GREEN}Detections: $detections${NC}"
        echo "$response" | jq '.detections[] | "\(.class_name): \(.confidence)"'
        return 0
    else
        echo -e "${RED}Inference failed${NC}"
        return 1
    fi
}

infer_batch() {
    local images=("$@")
    
    # Verify all images exist
    for image in "${images[@]}"; do
        if [ ! -f "$image" ]; then
            echo -e "${RED}Error: Image not found: $image${NC}"
            return 1
        fi
    done
    
    echo -e "${YELLOW}Running batch inference on ${#images[@]} images...${NC}"
    
    # Build JSON array
    json_array=$(printf '%s\n' "${images[@]}" | jq -R . | jq -s .)
    
    response=$(curl -s -m "$TIMEOUT" -X POST "$SERVICE_URL/api/v1/infer/batch" \
        -H "Content-Type: application/json" \
        -d "{\"image_paths\": $json_array}")
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}Batch inference completed${NC}"
        echo "$response" | jq '.'
        return 0
    else
        echo -e "${RED}Batch inference failed${NC}"
        return 1
    fi
}

# Main
if [ $# -eq 0 ]; then
    print_usage
    exit 0
fi

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --url)
            SERVICE_URL="$2"
            shift 2
            ;;
        --timeout)
            TIMEOUT="$2"
            shift 2
            ;;
        --help)
            print_usage
            exit 0
            ;;
        health)
            health_check
            exit $?
            ;;
        info)
            get_model_info
            exit $?
            ;;
        infer)
            shift
            infer_single "$@"
            exit $?
            ;;
        batch)
            shift
            infer_batch "$@"
            exit $?
            ;;
        *)
            echo -e "${RED}Unknown command: $1${NC}"
            print_usage
            exit 1
            ;;
    esac
done

